{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-04 22:49:03.765515: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-04 22:49:03.900707: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-03-04 22:49:04.449033: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/ubuntu/data/miniconda3/envs/tf2/lib/\n",
      "2023-03-04 22:49:04.449120: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/ubuntu/data/miniconda3/envs/tf2/lib/\n",
      "2023-03-04 22:49:04.449128: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_compression as tfc\n",
    "import tensorflow_probability as tfp\n",
    "\n",
    "from time import time\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import os\n",
    "import h5py\n",
    "\n",
    "from BalleFFP import BalleFFP\n",
    "from read_data import read_data_numpy\n",
    "\n",
    "import constants"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gpu_settings() -> tf.distribute.MirroredStrategy:\n",
    "    gpus = tf.config.list_physical_devices('GPU')\n",
    "    for gpu in gpus:\n",
    "        tf.config.experimental.set_memory_growth(gpu, True)\n",
    "    strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1']) # gpu distribution strategy\n",
    "    return strategy\n",
    "\n",
    "def read_data(ch_format='channels_first') -> np.ndarray:\n",
    "    data_path = os.path.join(constants.DATA_FOLDER, constants.DATA_FILE)\n",
    "    data      = read_data_numpy(data_path, ch_format)\n",
    "    return data\n",
    "\n",
    "def split_train_test(data: np.ndarray):\n",
    "    train_images = data[:constants.TRAINING_SET_SIZE]\n",
    "    test_images  = data[constants.TRAINING_SET_SIZE:constants.TRAINING_SET_SIZE+constants.VALIDATION_SET_SIZE]\n",
    "    return train_images, test_images"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-04 22:49:17.172045: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.173487: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.182052: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.183310: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.184588: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.185800: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.188239: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-03-04 22:49:17.564021: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.565492: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.566851: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.568470: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.569720: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:17.570930: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.575209: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.576773: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.578260: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.579477: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.580730: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.581961: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13641 MB memory:  -> device: 0, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5\n",
      "2023-03-04 22:49:18.582543: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-03-04 22:49:18.583819: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 13641 MB memory:  -> device: 1, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of devices: 2\n",
      "Reading data...\n",
      "Data shape: (100000, 3, 96, 96)\n"
     ]
    }
   ],
   "source": [
    "strategy  = gpu_settings()\n",
    "\n",
    "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))\n",
    "\n",
    "print('Reading data...')\n",
    "ch_format = 'channels_first'\n",
    "data      = read_data(ch_format)\n",
    "data      = data.astype('float32') # / 255.0\n",
    "print('Data shape: {}'.format(data.shape))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training/Test distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distributing data...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-04 22:49:51.037455: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:784] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: \"TensorSliceDataset/_1\"\n",
      "op: \"TensorSliceDataset\"\n",
      "input: \"Placeholder/_0\"\n",
      "attr {\n",
      "  key: \"Toutput_types\"\n",
      "  value {\n",
      "    list {\n",
      "      type: DT_FLOAT\n",
      "    }\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"_cardinality\"\n",
      "  value {\n",
      "    i: 80000\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"is_files\"\n",
      "  value {\n",
      "    b: false\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"metadata\"\n",
      "  value {\n",
      "    s: \"\\n\\024TensorSliceDataset:0\"\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"output_shapes\"\n",
      "  value {\n",
      "    list {\n",
      "      shape {\n",
      "        dim {\n",
      "          size: 3\n",
      "        }\n",
      "        dim {\n",
      "          size: 96\n",
      "        }\n",
      "        dim {\n",
      "          size: 96\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"replicate_on_split\"\n",
      "  value {\n",
      "    b: false\n",
      "  }\n",
      "}\n",
      "experimental_type {\n",
      "  type_id: TFT_PRODUCT\n",
      "  args {\n",
      "    type_id: TFT_DATASET\n",
      "    args {\n",
      "      type_id: TFT_PRODUCT\n",
      "      args {\n",
      "        type_id: TFT_TENSOR\n",
      "        args {\n",
      "          type_id: TFT_FLOAT\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n",
      "2023-03-04 22:49:52.745659: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:784] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: \"TensorSliceDataset/_1\"\n",
      "op: \"TensorSliceDataset\"\n",
      "input: \"Placeholder/_0\"\n",
      "attr {\n",
      "  key: \"Toutput_types\"\n",
      "  value {\n",
      "    list {\n",
      "      type: DT_FLOAT\n",
      "    }\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"_cardinality\"\n",
      "  value {\n",
      "    i: 20000\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"is_files\"\n",
      "  value {\n",
      "    b: false\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"metadata\"\n",
      "  value {\n",
      "    s: \"\\n\\024TensorSliceDataset:3\"\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"output_shapes\"\n",
      "  value {\n",
      "    list {\n",
      "      shape {\n",
      "        dim {\n",
      "          size: 3\n",
      "        }\n",
      "        dim {\n",
      "          size: 96\n",
      "        }\n",
      "        dim {\n",
      "          size: 96\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"replicate_on_split\"\n",
      "  value {\n",
      "    b: false\n",
      "  }\n",
      "}\n",
      "experimental_type {\n",
      "  type_id: TFT_PRODUCT\n",
      "  args {\n",
      "    type_id: TFT_DATASET\n",
      "    args {\n",
      "      type_id: TFT_PRODUCT\n",
      "      args {\n",
      "        type_id: TFT_TENSOR\n",
      "        args {\n",
      "          type_id: TFT_FLOAT\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train_images, test_images = split_train_test(data)\n",
    "del data\n",
    "\n",
    "buffer_size       = len(train_images) # buffer size for shuffling\n",
    "global_batch_size = constants.BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync # global batch size (in our case 2gpu * BATCH_SIZE_PER_REPLICA)\n",
    "\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(buffer_size).batch(global_batch_size)\n",
    "test_dataset  = tf.data.Dataset.from_tensor_slices(test_images).batch(global_batch_size)\n",
    "\n",
    "print('Distributing data...')\n",
    "train_dataset_dist = strategy.experimental_distribute_dataset(train_dataset)\n",
    "test_dataset_dist  = strategy.experimental_distribute_dataset(test_dataset)\n",
    "\n",
    "del train_dataset, test_dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    vae       = BalleFFP(N=128, M=192, k2=3, c=3, format=ch_format)\n",
    "    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4) "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training/Test steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    \n",
    "    @tf.function\n",
    "    def Loss(inputs, outputs):\n",
    "        return tf.reduce_mean(tf.square(inputs - outputs[0])) + 0.5*(tf.reduce_mean(outputs[1]))\n",
    "\n",
    "    @tf.function # compile the function to a graph for faster execution\n",
    "    def train_step(inputs, vae):\n",
    "        with tf.GradientTape() as tape: # create a tape to record operations\n",
    "            reconstructed, rateb = vae(inputs) # forward pass\n",
    "            loss = Loss(inputs, (reconstructed, rateb)) # MSE loss (maybe put this in a function)\n",
    "        gradients = tape.gradient(loss, vae.trainable_variables) # compute gradients    \n",
    "        optimizer.apply_gradients(zip(gradients, vae.trainable_variables)) # gradient descent\n",
    "        return loss # return loss for logging\n",
    "    \n",
    "    @tf.function\n",
    "    def val_step(inputs, vae):\n",
    "        outputs = vae(inputs, training=False) # forward pass\n",
    "        loss    = Loss(inputs, outputs)       \n",
    "        return loss \n",
    "\n",
    "\n",
    "def train_step_dist(inputs, vae, strategy):\n",
    "    loss = strategy.run(train_step, args=(inputs, vae))\n",
    "    return strategy.reduce(tf.distribute.ReduceOp.SUM, loss, axis=None)\n",
    "\n",
    "def val_step_dist(inputs, vae, strategy):\n",
    "    loss = strategy.run(val_step, args=(inputs, vae))\n",
    "    return strategy.reduce(tf.distribute.ReduceOp.SUM, loss, axis=None)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training started...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1faac1dcd2904536ad084c87eee99f3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "training steps: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/ubuntu/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
      "Instructions for updating:\n",
      "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089\n",
      "INFO:tensorflow:batch_all_reduce: 36 all-reduces with algorithm = nccl, num_packs = 1\n",
      "INFO:tensorflow:batch_all_reduce: 36 all-reduces with algorithm = nccl, num_packs = 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-04 22:50:12.468987: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-03-04 22:50:13.129854: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:428] Loaded cuDNN version 8100\n",
      "2023-03-04 22:50:14.784216: I tensorflow/compiler/xla/service/service.cc:173] XLA service 0x7f5370975970 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:\n",
      "2023-03-04 22:50:14.784257: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): Tesla T4, Compute Capability 7.5\n",
      "2023-03-04 22:50:14.784262: I tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (1): Tesla T4, Compute Capability 7.5\n",
      "2023-03-04 22:50:14.790711: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
      "2023-03-04 22:50:14.917299: I tensorflow/compiler/jit/xla_compilation_cache.cc:477] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n",
      "INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 9\u001b[0m\n\u001b[1;32m      7\u001b[0m total_loss  \u001b[39m=\u001b[39m \u001b[39m0.0\u001b[39m \n\u001b[1;32m      8\u001b[0m num_batches \u001b[39m=\u001b[39m \u001b[39m0\u001b[39m \n\u001b[0;32m----> 9\u001b[0m \u001b[39mfor\u001b[39;00m inputs \u001b[39min\u001b[39;00m tqdm(train_dataset_dist, \u001b[39m'\u001b[39m\u001b[39mtraining steps\u001b[39m\u001b[39m'\u001b[39m): \n\u001b[1;32m     10\u001b[0m     total_loss \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m train_step_dist(inputs, vae, strategy) \u001b[39m# type: ignore # sum losses across replicas (type: ignore is for mypy\u001b[39;00m\n\u001b[1;32m     11\u001b[0m     num_batches \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m \u001b[39m# count number of batches\u001b[39;00m\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tqdm/notebook.py:259\u001b[0m, in \u001b[0;36mtqdm_notebook.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    257\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m    258\u001b[0m     it \u001b[39m=\u001b[39m \u001b[39msuper\u001b[39m(tqdm_notebook, \u001b[39mself\u001b[39m)\u001b[39m.\u001b[39m\u001b[39m__iter__\u001b[39m()\n\u001b[0;32m--> 259\u001b[0m     \u001b[39mfor\u001b[39;00m obj \u001b[39min\u001b[39;00m it:\n\u001b[1;32m    260\u001b[0m         \u001b[39m# return super(tqdm...) will not catch exception\u001b[39;00m\n\u001b[1;32m    261\u001b[0m         \u001b[39myield\u001b[39;00m obj\n\u001b[1;32m    262\u001b[0m \u001b[39m# NB: except ... [ as ...] breaks IPython async KeyboardInterrupt\u001b[39;00m\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tqdm/std.py:1195\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1192\u001b[0m time \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_time\n\u001b[1;32m   1194\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1195\u001b[0m     \u001b[39mfor\u001b[39;00m obj \u001b[39min\u001b[39;00m iterable:\n\u001b[1;32m   1196\u001b[0m         \u001b[39myield\u001b[39;00m obj\n\u001b[1;32m   1197\u001b[0m         \u001b[39m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   1198\u001b[0m         \u001b[39m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/distribute/input_lib.py:573\u001b[0m, in \u001b[0;36mDistributedIteratorBase.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    571\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__next__\u001b[39m(\u001b[39mself\u001b[39m):\n\u001b[1;32m    572\u001b[0m   \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 573\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mget_next()\n\u001b[1;32m    574\u001b[0m   \u001b[39mexcept\u001b[39;00m errors\u001b[39m.\u001b[39mOutOfRangeError:\n\u001b[1;32m    575\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mStopIteration\u001b[39;00m\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/distribute/input_lib.py:636\u001b[0m, in \u001b[0;36mDistributedIteratorBase.get_next\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m    634\u001b[0m   \u001b[39mwith\u001b[39;00m ops\u001b[39m.\u001b[39mdevice(worker):\n\u001b[1;32m    635\u001b[0m     optional_list\u001b[39m.\u001b[39mappend(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_iterators[i]\u001b[39m.\u001b[39mget_next_as_optional_list())\n\u001b[0;32m--> 636\u001b[0m num_replicas_with_values \u001b[39m=\u001b[39m _calculate_replicas_with_values(\n\u001b[1;32m    637\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_strategy, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_input_workers, optional_list)\n\u001b[1;32m    639\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_value_or_dummy\u001b[39m():\n\u001b[1;32m    640\u001b[0m   value_list \u001b[39m=\u001b[39m _get_value_or_dummy(\n\u001b[1;32m    641\u001b[0m       \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_input_workers, optional_list, produce_dummy\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/distribute/input_lib.py:514\u001b[0m, in \u001b[0;36m_calculate_replicas_with_values\u001b[0;34m(strategy, input_workers, optional_list)\u001b[0m\n\u001b[1;32m    509\u001b[0m     device_has_values \u001b[39m=\u001b[39m [\n\u001b[1;32m    510\u001b[0m         math_ops\u001b[39m.\u001b[39mcast(v\u001b[39m.\u001b[39mhas_value(), dtypes\u001b[39m.\u001b[39mint64) \u001b[39mfor\u001b[39;00m v \u001b[39min\u001b[39;00m optionals\n\u001b[1;32m    511\u001b[0m     ]\n\u001b[1;32m    512\u001b[0m     worker_has_values\u001b[39m.\u001b[39mappend(\n\u001b[1;32m    513\u001b[0m         math_ops\u001b[39m.\u001b[39mreduce_sum(device_has_values, keepdims\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m))\n\u001b[0;32m--> 514\u001b[0m client_has_values \u001b[39m=\u001b[39m math_ops\u001b[39m.\u001b[39;49mreduce_sum(worker_has_values, keepdims\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m    515\u001b[0m \u001b[39mif\u001b[39;00m strategy\u001b[39m.\u001b[39mextended\u001b[39m.\u001b[39m_in_multi_worker_mode():  \u001b[39m# pylint: disable=protected-access\u001b[39;00m\n\u001b[1;32m    516\u001b[0m   global_has_values \u001b[39m=\u001b[39m strategy\u001b[39m.\u001b[39mreduce(\n\u001b[1;32m    517\u001b[0m       reduce_util\u001b[39m.\u001b[39mReduceOp\u001b[39m.\u001b[39mSUM, client_has_values, axis\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:150\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    148\u001b[0m filtered_tb \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    149\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 150\u001b[0m   \u001b[39mreturn\u001b[39;00m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    151\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    152\u001b[0m   filtered_tb \u001b[39m=\u001b[39m _process_traceback_frames(e\u001b[39m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py:1176\u001b[0m, in \u001b[0;36madd_dispatch_support.<locals>.decorator.<locals>.op_dispatch_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1174\u001b[0m \u001b[39m# Fallback dispatch system (dispatch v1):\u001b[39;00m\n\u001b[1;32m   1175\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1176\u001b[0m   \u001b[39mreturn\u001b[39;00m dispatch_target(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1177\u001b[0m \u001b[39mexcept\u001b[39;00m (\u001b[39mTypeError\u001b[39;00m, \u001b[39mValueError\u001b[39;00m):\n\u001b[1;32m   1178\u001b[0m   \u001b[39m# Note: convert_to_eager_tensor currently raises a ValueError, not a\u001b[39;00m\n\u001b[1;32m   1179\u001b[0m   \u001b[39m# TypeError, when given unexpected types.  So we need to catch both.\u001b[39;00m\n\u001b[1;32m   1180\u001b[0m   result \u001b[39m=\u001b[39m dispatch(op_dispatch_handler, args, kwargs)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py:2313\u001b[0m, in \u001b[0;36mreduce_sum\u001b[0;34m(input_tensor, axis, keepdims, name)\u001b[0m\n\u001b[1;32m   2249\u001b[0m \u001b[39m@tf_export\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mmath.reduce_sum\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mreduce_sum\u001b[39m\u001b[39m\"\u001b[39m, v1\u001b[39m=\u001b[39m[])\n\u001b[1;32m   2250\u001b[0m \u001b[39m@dispatch\u001b[39m\u001b[39m.\u001b[39madd_dispatch_support\n\u001b[1;32m   2251\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mreduce_sum\u001b[39m(input_tensor, axis\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, keepdims\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m, name\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m):\n\u001b[1;32m   2252\u001b[0m \u001b[39m  \u001b[39m\u001b[39m\"\"\"Computes the sum of elements across dimensions of a tensor.\u001b[39;00m\n\u001b[1;32m   2253\u001b[0m \n\u001b[1;32m   2254\u001b[0m \u001b[39m  This is the reduction operation for the elementwise `tf.math.add` op.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   2309\u001b[0m \u001b[39m  @end_compatibility\u001b[39;00m\n\u001b[1;32m   2310\u001b[0m \u001b[39m  \"\"\"\u001b[39;00m\n\u001b[1;32m   2312\u001b[0m   \u001b[39mreturn\u001b[39;00m reduce_sum_with_dims(input_tensor, axis, keepdims, name,\n\u001b[0;32m-> 2313\u001b[0m                               _ReductionDims(input_tensor, axis))\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py:2153\u001b[0m, in \u001b[0;36m_ReductionDims\u001b[0;34m(x, axis)\u001b[0m\n\u001b[1;32m   2150\u001b[0m   \u001b[39mreturn\u001b[39;00m constant_op\u001b[39m.\u001b[39mconstant(np\u001b[39m.\u001b[39marange(x_rank, dtype\u001b[39m=\u001b[39mnp\u001b[39m.\u001b[39mint32))\n\u001b[1;32m   2151\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   2152\u001b[0m   \u001b[39m# Otherwise, we rely on Range and Rank to do the right thing at run-time.\u001b[39;00m\n\u001b[0;32m-> 2153\u001b[0m   \u001b[39mreturn\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39m0\u001b[39m, array_ops\u001b[39m.\u001b[39;49mrank(x))\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:150\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    148\u001b[0m filtered_tb \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    149\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 150\u001b[0m   \u001b[39mreturn\u001b[39;00m fn(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m    151\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    152\u001b[0m   filtered_tb \u001b[39m=\u001b[39m _process_traceback_frames(e\u001b[39m.\u001b[39m__traceback__)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py:1176\u001b[0m, in \u001b[0;36madd_dispatch_support.<locals>.decorator.<locals>.op_dispatch_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1174\u001b[0m \u001b[39m# Fallback dispatch system (dispatch v1):\u001b[39;00m\n\u001b[1;32m   1175\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 1176\u001b[0m   \u001b[39mreturn\u001b[39;00m dispatch_target(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1177\u001b[0m \u001b[39mexcept\u001b[39;00m (\u001b[39mTypeError\u001b[39;00m, \u001b[39mValueError\u001b[39;00m):\n\u001b[1;32m   1178\u001b[0m   \u001b[39m# Note: convert_to_eager_tensor currently raises a ValueError, not a\u001b[39;00m\n\u001b[1;32m   1179\u001b[0m   \u001b[39m# TypeError, when given unexpected types.  So we need to catch both.\u001b[39;00m\n\u001b[1;32m   1180\u001b[0m   result \u001b[39m=\u001b[39m dispatch(op_dispatch_handler, args, kwargs)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py:857\u001b[0m, in \u001b[0;36mrank\u001b[0;34m(input, name)\u001b[0m\n\u001b[1;32m    824\u001b[0m \u001b[39m@tf_export\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mrank\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m    825\u001b[0m \u001b[39m@dispatch\u001b[39m\u001b[39m.\u001b[39madd_dispatch_support\n\u001b[1;32m    826\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mrank\u001b[39m(\u001b[39minput\u001b[39m, name\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m):\n\u001b[1;32m    827\u001b[0m   \u001b[39m# pylint: disable=redefined-builtin\u001b[39;00m\n\u001b[1;32m    828\u001b[0m \u001b[39m  \u001b[39m\u001b[39m\"\"\"Returns the rank of a tensor.\u001b[39;00m\n\u001b[1;32m    829\u001b[0m \n\u001b[1;32m    830\u001b[0m \u001b[39m  See also `tf.shape`.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    855\u001b[0m \u001b[39m  @end_compatibility\u001b[39;00m\n\u001b[1;32m    856\u001b[0m \u001b[39m  \"\"\"\u001b[39;00m\n\u001b[0;32m--> 857\u001b[0m   \u001b[39mreturn\u001b[39;00m rank_internal(\u001b[39minput\u001b[39;49m, name, optimize\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py:877\u001b[0m, in \u001b[0;36mrank_internal\u001b[0;34m(input, name, optimize)\u001b[0m\n\u001b[1;32m    875\u001b[0m   \u001b[39mreturn\u001b[39;00m gen_array_ops\u001b[39m.\u001b[39msize(\u001b[39minput\u001b[39m\u001b[39m.\u001b[39mdense_shape, name\u001b[39m=\u001b[39mname)\n\u001b[1;32m    876\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m--> 877\u001b[0m   \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m ops\u001b[39m.\u001b[39;49mconvert_to_tensor(\u001b[39minput\u001b[39;49m)\n\u001b[1;32m    878\u001b[0m   input_shape \u001b[39m=\u001b[39m \u001b[39minput\u001b[39m\u001b[39m.\u001b[39mget_shape()\n\u001b[1;32m    879\u001b[0m   \u001b[39mif\u001b[39;00m optimize \u001b[39mand\u001b[39;00m input_shape\u001b[39m.\u001b[39mndims \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/profiler/trace.py:183\u001b[0m, in \u001b[0;36mtrace_wrapper.<locals>.inner_wrapper.<locals>.wrapped\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    181\u001b[0m   \u001b[39mwith\u001b[39;00m Trace(trace_name, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mtrace_kwargs):\n\u001b[1;32m    182\u001b[0m     \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m--> 183\u001b[0m \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/framework/ops.py:1636\u001b[0m, in \u001b[0;36mconvert_to_tensor\u001b[0;34m(value, dtype, name, as_ref, preferred_dtype, dtype_hint, ctx, accepted_result_types)\u001b[0m\n\u001b[1;32m   1627\u001b[0m       \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[1;32m   1628\u001b[0m           _add_error_prefix(\n\u001b[1;32m   1629\u001b[0m               \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mConversion function \u001b[39m\u001b[39m{\u001b[39;00mconversion_func\u001b[39m!r}\u001b[39;00m\u001b[39m for type \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1632\u001b[0m               \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mactual = \u001b[39m\u001b[39m{\u001b[39;00mret\u001b[39m.\u001b[39mdtype\u001b[39m.\u001b[39mbase_dtype\u001b[39m.\u001b[39mname\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m,\n\u001b[1;32m   1633\u001b[0m               name\u001b[39m=\u001b[39mname))\n\u001b[1;32m   1635\u001b[0m \u001b[39mif\u001b[39;00m ret \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[0;32m-> 1636\u001b[0m   ret \u001b[39m=\u001b[39m conversion_func(value, dtype\u001b[39m=\u001b[39;49mdtype, name\u001b[39m=\u001b[39;49mname, as_ref\u001b[39m=\u001b[39;49mas_ref)\n\u001b[1;32m   1638\u001b[0m \u001b[39mif\u001b[39;00m ret \u001b[39mis\u001b[39;00m \u001b[39mNotImplemented\u001b[39m:\n\u001b[1;32m   1639\u001b[0m   \u001b[39mcontinue\u001b[39;00m\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py:1589\u001b[0m, in \u001b[0;36m_autopacking_conversion_function\u001b[0;34m(v, dtype, name, as_ref)\u001b[0m\n\u001b[1;32m   1587\u001b[0m \u001b[39melif\u001b[39;00m dtype \u001b[39m!=\u001b[39m inferred_dtype:\n\u001b[1;32m   1588\u001b[0m   v \u001b[39m=\u001b[39m nest\u001b[39m.\u001b[39mmap_structure(_cast_nested_seqs_to_dtype(dtype), v)\n\u001b[0;32m-> 1589\u001b[0m \u001b[39mreturn\u001b[39;00m _autopacking_helper(v, dtype, name \u001b[39mor\u001b[39;49;00m \u001b[39m\"\u001b[39;49m\u001b[39mpacked\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/array_ops.py:1496\u001b[0m, in \u001b[0;36m_autopacking_helper\u001b[0;34m(list_or_tuple, dtype, name)\u001b[0m\n\u001b[1;32m   1492\u001b[0m \u001b[39mif\u001b[39;00m context\u001b[39m.\u001b[39mexecuting_eagerly():\n\u001b[1;32m   1493\u001b[0m   \u001b[39m# NOTE: Fast path when all the items are tensors, this doesn't do any type\u001b[39;00m\n\u001b[1;32m   1494\u001b[0m   \u001b[39m# checking.\u001b[39;00m\n\u001b[1;32m   1495\u001b[0m   \u001b[39mif\u001b[39;00m \u001b[39mall\u001b[39m(\u001b[39misinstance\u001b[39m(elem, core\u001b[39m.\u001b[39mTensor) \u001b[39mfor\u001b[39;00m elem \u001b[39min\u001b[39;00m list_or_tuple):\n\u001b[0;32m-> 1496\u001b[0m     \u001b[39mreturn\u001b[39;00m gen_array_ops\u001b[39m.\u001b[39;49mpack(list_or_tuple, name\u001b[39m=\u001b[39;49mname)\n\u001b[1;32m   1497\u001b[0m must_pack \u001b[39m=\u001b[39m \u001b[39mFalse\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m converted_elems \u001b[39m=\u001b[39m []\n",
      "File \u001b[0;32m~/data/miniconda3/envs/tf2/lib/python3.9/site-packages/tensorflow/python/ops/gen_array_ops.py:6549\u001b[0m, in \u001b[0;36mpack\u001b[0;34m(values, axis, name)\u001b[0m\n\u001b[1;32m   6547\u001b[0m \u001b[39mif\u001b[39;00m tld\u001b[39m.\u001b[39mis_eager:\n\u001b[1;32m   6548\u001b[0m   \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> 6549\u001b[0m     _result \u001b[39m=\u001b[39m pywrap_tfe\u001b[39m.\u001b[39;49mTFE_Py_FastPathExecute(\n\u001b[1;32m   6550\u001b[0m       _ctx, \u001b[39m\"\u001b[39;49m\u001b[39mPack\u001b[39;49m\u001b[39m\"\u001b[39;49m, name, values, \u001b[39m\"\u001b[39;49m\u001b[39maxis\u001b[39;49m\u001b[39m\"\u001b[39;49m, axis)\n\u001b[1;32m   6551\u001b[0m     \u001b[39mreturn\u001b[39;00m _result\n\u001b[1;32m   6552\u001b[0m   \u001b[39mexcept\u001b[39;00m _core\u001b[39m.\u001b[39m_NotOkStatusException \u001b[39mas\u001b[39;00m e:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "training_losses = []\n",
    "test_losses     = []\n",
    "\n",
    "print('Training started...')\n",
    "\n",
    "for epoch in range(constants.EPOCHS):\n",
    "    total_loss  = 0.0 \n",
    "    num_batches = 0 \n",
    "    for inputs in tqdm(train_dataset_dist, 'training steps'): \n",
    "        total_loss += train_step_dist(inputs, vae, strategy) # type: ignore # sum losses across replicas (type: ignore is for mypy\n",
    "        num_batches += 1 # count number of batches\n",
    "        \n",
    "    train_loss = total_loss / num_batches # compute average loss\n",
    "    training_losses.append(train_loss)\n",
    "    \n",
    "    print('Epoch {} train loss: {}'.format(epoch, train_loss))\n",
    "\n",
    "    total_loss  = 0.0\n",
    "    num_batches = 0\n",
    "    for inputs in tqdm(test_dataset_dist, 'validation steps'): \n",
    "        total_loss  += val_step_dist(inputs, vae, strategy) # type: ignore # sum losses across replicas\n",
    "        num_batches += 1 # count number of batches\n",
    "        \n",
    "    test_loss = total_loss / num_batches # compute average loss\n",
    "    test_losses.append(test_loss)\n",
    "    \n",
    "    print('Epoch {} test loss: {}'.format(epoch, test_loss))\n",
    "    \n",
    "print('Training finished!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Saving model and losses...')\n",
    "current_time = time()\n",
    "\n",
    "# save model\n",
    "model_name = f\"model_{current_time}.h5\"\n",
    "model_path = constants.MODEL_FOLDER + model_name\n",
    "vae.save_weights(model_path)\n",
    "\n",
    "# save losses in .h5 file\n",
    "losses_name = f\"losses_{current_time}.h5\"\n",
    "losses_path = constants.MODEL_FOLDER + losses_name\n",
    "with h5py.File(losses_path, 'w') as f:\n",
    "    f.create_dataset('train', data=training_losses)\n",
    "    f.create_dataset('test', data=test_losses)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
